{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Runtime Logging with AG2 \n",
    "\n",
    "AG2 offers utilities to log data for debugging and performance analysis. This notebook demonstrates how to use them. \n",
    "\n",
    "we log data in different modes:\n",
    "- SQlite Database\n",
    "- File \n",
    "\n",
    "In general, users can initiate logging by calling `autogen.runtime_logging.start()` and stop logging by calling `autogen.runtime_logging.stop()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent, LLMConfig, UserProxyAgent\n",
    "\n",
    "# Setup API key. Add your own API key to config file or environment variable\n",
    "llm_config = LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", temperature=0.9)\n",
    "\n",
    "# Start logging\n",
    "logging_session_id = autogen.runtime_logging.start(config={\"dbname\": \"logs.db\"})\n",
    "print(\"Logging session ID: \" + str(logging_session_id))\n",
    "\n",
    "# Create an agent workflow and run it\n",
    "assistant = AssistantAgent(name=\"assistant\", llm_config=llm_config)\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    code_execution_config=False,\n",
    "    human_input_mode=\"NEVER\",\n",
    "    is_termination_msg=lambda msg: \"TERMINATE\" in msg[\"content\"],\n",
    ")\n",
    "\n",
    "user_proxy.initiate_chat(\n",
    "    assistant, message=\"What is the height of the Eiffel Tower? Only respond with the answer and terminate\"\n",
    ")\n",
    "autogen.runtime_logging.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting Data from the SQLite Database \n",
    "\n",
    "`logs.db` should be generated, by default it's using SQLite database. You can view the data with GUI tool like `sqlitebrowser`, using SQLite command line shell or using python script:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_log(dbname=\"logs.db\", table=\"chat_completions\"):\n",
    "    import sqlite3\n",
    "\n",
    "    con = sqlite3.connect(dbname)\n",
    "    query = f\"SELECT * from {table}\"\n",
    "    cursor = con.execute(query)\n",
    "    rows = cursor.fetchall()\n",
    "    column_names = [description[0] for description in cursor.description]\n",
    "    data = [dict(zip(column_names, row)) for row in rows]\n",
    "    con.close()\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def str_to_dict(s):\n",
    "    return json.loads(s)\n",
    "\n",
    "\n",
    "log_data = get_log()\n",
    "log_data_df = pd.DataFrame(log_data)\n",
    "\n",
    "log_data_df[\"total_tokens\"] = log_data_df.apply(\n",
    "    lambda row: str_to_dict(row[\"response\"])[\"usage\"][\"total_tokens\"], axis=1\n",
    ")\n",
    "\n",
    "log_data_df[\"request\"] = log_data_df.apply(lambda row: str_to_dict(row[\"request\"])[\"messages\"][0][\"content\"], axis=1)\n",
    "\n",
    "log_data_df[\"response\"] = log_data_df.apply(\n",
    "    lambda row: str_to_dict(row[\"response\"])[\"choices\"][0][\"message\"][\"content\"], axis=1\n",
    ")\n",
    "\n",
    "log_data_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing Cost \n",
    "\n",
    "One use case of logging data is to compute the cost of a session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sum totoal tokens for all sessions\n",
    "total_tokens = log_data_df[\"total_tokens\"].sum()\n",
    "\n",
    "# Sum total cost for all sessions\n",
    "total_cost = log_data_df[\"cost\"].sum()\n",
    "\n",
    "# Total tokens for specific session\n",
    "session_tokens = log_data_df[log_data_df[\"session_id\"] == logging_session_id][\"total_tokens\"].sum()\n",
    "session_cost = log_data_df[log_data_df[\"session_id\"] == logging_session_id][\"cost\"].sum()\n",
    "\n",
    "print(\"Total tokens for all sessions: \" + str(total_tokens) + \", total cost: \" + str(round(total_cost, 4)))\n",
    "print(\n",
    "    \"Total tokens for session \"\n",
    "    + str(logging_session_id)\n",
    "    + \": \"\n",
    "    + str(session_tokens)\n",
    "    + \", cost: \"\n",
    "    + str(round(session_cost, 4))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log data in File mode\n",
    "\n",
    "By default, the log type is set to `sqlite` as shown above, but we introduced a new parameter for the `autogen.runtime_logging.start()`\n",
    "\n",
    "the `logger_type = \"file\"` will start to log data in the File mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "import autogen\n",
    "from autogen import AssistantAgent, LLMConfig, UserProxyAgent\n",
    "\n",
    "# Setup API key. Add your own API key to config file or environment variable\n",
    "llm_config = LLMConfig.from_json(path=\"OAI_CONFIG_LIST\", temperature=0.9)\n",
    "\n",
    "# Start logging with logger_type and the filename to log to\n",
    "logging_session_id = autogen.runtime_logging.start(logger_type=\"file\", config={\"filename\": \"runtime.log\"})\n",
    "print(\"Logging session ID: \" + str(logging_session_id))\n",
    "\n",
    "# Create an agent workflow and run it\n",
    "assistant = AssistantAgent(name=\"assistant\", llm_config=llm_config)\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    code_execution_config=False,\n",
    "    human_input_mode=\"NEVER\",\n",
    "    is_termination_msg=lambda msg: \"TERMINATE\" in msg[\"content\"],\n",
    ")\n",
    "\n",
    "user_proxy.initiate_chat(\n",
    "    assistant, message=\"What is the height of the Eiffel Tower? Only respond with the answer and terminate\"\n",
    ")\n",
    "autogen.runtime_logging.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This should create a `runtime.log` file in your current directory. "
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Provide capabilities of runtime logging for debugging and performance analysis.",
   "tags": [
    "logging",
    "debugging"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
